from __future__ import unicode_literals
from django.db import models
from django.utils.encoding import python_2_unicode_compatible
from django_rdkit.models import *
from django_rdkit.models.fields import *
from django_rdkit.config import config
from rdkit import Chem
from rdkit.Chem import Descriptors
from rdkit.Chem.Crippen import MolLogP
from rdkit.Chem.Lipinski import NumHAcceptors, NumHDonors, NumRotatableBonds

class CompoundManager(models.Manager):

    def structure_search(self, smiles, similarity):
        search_mfp2 = MORGANBV_FP(Value(smiles))
        config.tanimoto_threshold = similarity
        queryset = super(CompoundManager, self).get_queryset().filter(mfp2__tanimoto=search_mfp2)
        queryset = queryset.annotate(similarity=TANIMOTO_SML('mfp2', search_mfp2))
        queryset = queryset.order_by('-similarity')
        return queryset


@python_2_unicode_compatible
class Compound(models.Model):
    objects = CompoundManager()
    generic_name = models.CharField(max_length=200, blank=True)
    IUPAC_name = models.CharField(max_length=2048, blank=True)
    smiles = models.CharField(max_length=2048, blank=True)
    image = models.ImageField(upload_to='mol_images', blank=True, null=True)
    mol_weight = models.FloatField(null=True, blank=True)
    cas = models.CharField(max_length=200, blank=True, null=True)
    cid = models.CharField(max_length=200, blank=True, null=True)
    drugbank_id = models.CharField(max_length=200, blank=True, null=True)
    #drug_groups = models.CharField(max_length=200, blank=True, null=True)

    first_created = models.DateTimeField(auto_now_add=True)
    last_created = models.DateTimeField(auto_now_add=True)

    drug_status = models.CharField(max_length=1024, blank=True, null=True)
    bindingdb_link = models.URLField(blank=True)
    chembl_link = models.URLField(blank=True)
    pathway_link = models.URLField(blank=True)

    mol = MolField(null=True, blank=True)
    torsionbv = BfpField(null=True)
    mfp2 = BfpField(null=True)
    ffp2 = BfpField(null=True)
    mol_block = models.TextField(blank=True)
    formula = models.CharField(max_length=1024, blank=True)
    alogp = models.DecimalField(
        max_digits=9,
        decimal_places=2,
        verbose_name=('Calculated AlogP'),
        blank=True,
        null=True
    )
    hba = models.SmallIntegerField(
        verbose_name=('Number hydrogen bond acceptors'),
        blank=True,
        null=True
    )
    hbd = models.SmallIntegerField(
        verbose_name=('Number hydrogen bond donors'),
        blank=True,
        null=True
    )
    psa = models.DecimalField(
        max_digits=9,
        decimal_places=2,
        verbose_name=('Polar surface area'),
        blank=True,
        null=True
    )
    rtb = models.SmallIntegerField(
        verbose_name=('Number rotabable bonds'),
        blank=True,
        null=True
    )

    def save(self, force_insert=False, force_update=False, using=None,
             update_fields=None, *args, **kwargs):
        smiles = self.smiles
        if smiles:
            try:
                self.mol = Chem.MolFromSmiles(smiles)
                self.mol_block = Chem.MolToMolBlock(self.mol)
                self.mol_weight = Descriptors.ExactMolWt(self.mol)
                self.alogp = MolLogP(self.mol)
                self.hba = NumHAcceptors(self.mol)
                self.hbd = NumHDonors(self.mol)
                self.psa = Chem.MolSurf.TPSA(self.mol)
                self.rtb = NumRotatableBonds(self.mol)
                super(Compound, self).save(*args, **kwargs)
                self.formula = Chem.rdMolDescriptors.CalcMolFormula(self.mol)
                #self.mfp2 = MORGANBV_FP(self.mol)
            except (ValueError, TypeError):
                print "Error when storing mol object"
        super(Compound, self).save(*args, **kwargs)

    def __str__(self):
        return '{}'.format(self.generic_name or self.smiles or self.pk)

@python_2_unicode_compatible
class Synonyms(models.Model):
    name = models.CharField(max_length=1024, blank=True, null=True)
    compound = models.ForeignKey('Compound', null=True, blank=True)

    def __str__(self):
        return self.name

@python_2_unicode_compatible
class Products(models.Model):
    name = models.CharField(max_length=1024, blank=True, null=True)
    compound = models.ForeignKey('Compound', null=True, blank=True)

    def __str__(self):
        return self.name

# @python_2_unicode_compatible
# class UniprotEntry(models.Model):
#     name = models.CharField(max_length=200, blank=True)
#     compounds = models.ManyToManyField('Compound', blank=True)
#
#     def __str__(self):
#         return self.name

@python_2_unicode_compatible
class UniprotInfo(models.Model):
    entry = models.CharField(max_length=200, blank=True, null=True)
    entryname = models.CharField(max_length=200, blank=True, null=True)
    compounds = models.ManyToManyField('Compound', blank=True)

    def __str__(self):
        return self.entry




